package nutcore

import chisel3._
import chisel3.util._
import chisel3.util.experimental._
import chiseltest._
import chiseltest.formal._
import chiseltest.experimental._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}


import chiseltest.formal._
import chiseltest.experimental._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}

import org.scalatest.Tag
import org.scalatest.flatspec.AnyFlatSpec

import chiseltest.ChiselScalatestTester
import chiseltest.formal.backends.FormalEngineAnnotation

import utils._
import bus.simplebus._
import bus.axi4._
import device._

// Available solver engine annotations:
// 
// Z3EngineAnnotation
// CVC4EngineAnnotation
// BtormcEngineAnnotation
// 
// It's been tested that Z3EngineAnnotation cannot work with VecInit. You should
// use BtormcEngineAnnotation in most of circumstances.

class FormalCacheModSpec extends AnyFlatSpec with ChiselScalatestTester with Formal {
   
  "Cache Level 1" should "pass" in {
    verify(new SingleLevelCacheSpecExternal, Seq(BoundedCheck(20), BtormcEngineAnnotation))
  }

  // "Double Level Cache" should "pass" in {
  //   verify(new DoubleLevelCacheSpecExternal, Seq(BoundedCheck(30), BtormcEngineAnnotation))
  // }
}

class AXI4RAM[T <: AXI4Lite](_type: T = new AXI4, memByte: Int,
  useBlackBox: Boolean = false) extends AXI4SlaveModule(_type) with HasNutCoreParameter {

  val offsetBits = log2Up(memByte)
  val offsetMask = (1 << offsetBits) - 1
  def index(addr: UInt) = (addr & offsetMask.U) >> log2Ceil(DataBytes)
  def inRange(idx: UInt) = idx < (memByte / 8).U

  val wIdx = index(waddr) + writeBeatCnt
  val rIdx = index(raddr) + readBeatCnt
  val wen = in.w.fire && inRange(wIdx)

  val rdata = {
    val mem = Mem(memByte / DataBytes, Vec(DataBytes, UInt(8.W)))

    val wdata = VecInit.tabulate(DataBytes) { i => in.w.bits.data(8*(i+1)-1, 8*i) }
    when (wen) { mem.write(wIdx, wdata, in.w.bits.strb.asBools) }

    Cat(mem.read(rIdx).reverse)
  }

  in.r.bits.data := RegEnable(rdata, ren)
}

trait HasPrefetcherParameter extends HasNutCoreParameter {
  val supportPrefetch = HasDcache
}

// ## The basic idea of prefetcher:
// 
// the register prefetchReq always holds a prefetch request, but the address varies 
// along with the input request address.
// 
// getNewReq controls whether to use the prefetch request.
// - when input on fire, and input request is burst, and the address of current req
//   differs from the last (latched) one, which indicates a miss, getNewReq is set.
// - when requested address is MMIO, or the transmission is done, getNewReq is reset.
//
// The effect of prefetch:
//
// The prefetch command is basically same as a read. The prefetch controls only the 
// resp valid of cache
//
// Thus, the prefetcher is expected to:
// - being transparent when the address of two consecutive requests are same.
// - switch to prefetch request right after a miss
// - switch back right after a prefetch is sent.

trait HasCoherenceParameter extends HasNutCoreParameter {
  val supportCoh = HasDcache
}

class CoherenceManager extends Module with HasCoherenceParameter {
  val io = IO(new Bundle {
    val in = Flipped(new SimpleBusUC)
    val out = new Bundle {
      val mem = new SimpleBusUC
      val coh = new SimpleBusUC
    }
  })

  // state transition:
  // write: s_idle -> s_memWriteResp -> s_idle
  // read: s_idle -> s_probeResp -> (hit) s_probeForward -> s_idle
  //                             +> (miss) s_memReadReq -> s_memReadResp -> s_idle

  val s_idle :: s_probeResp :: s_probeForward :: s_memReadReq :: s_memReadResp :: s_memWriteResp :: Nil = Enum(6)
  val state = RegInit(s_idle)
  val inflight = state =/= s_idle

  val thisReq = io.in.req
  assert(!(thisReq.valid && !thisReq.bits.isRead() && !thisReq.bits.isWrite()))

  // when read, we should first probe dcache
  val reqLatch = RegEnable(thisReq.bits, !inflight && thisReq.bits.isRead())
  io.out.coh match { case c => {
    c.req.bits := thisReq.bits
    c.req.bits.cmd := SimpleBusCmd.probe
    c.resp.ready := true.B
  }}

  io.out.mem.req.bits := thisReq.bits
  // bind correct valid and ready signals
  io.out.mem.req.valid := false.B
  thisReq.ready := false.B
  io.out.coh.req.valid := false.B
  when (if (supportCoh) thisReq.bits.isWrite() else true.B) {
    io.out.mem.req.valid := thisReq.valid && !inflight
    thisReq.ready := io.out.mem.req.ready && !inflight
  } .elsewhen (thisReq.bits.isRead()) {
    io.out.coh.req.valid := thisReq.valid && !inflight
    thisReq.ready := io.out.coh.req.ready && !inflight
  }

  io.in.resp <> io.out.mem.resp

  switch (state) {
    is (s_idle) {
      when (thisReq.fire()) {
        when (thisReq.bits.isRead()) { state := Mux(supportCoh.B, s_probeResp, s_memReadResp) }
        .elsewhen (thisReq.bits.isWriteLast()) { state := s_memWriteResp }
      }
    }
    is (s_probeResp) {
      when (io.out.coh.resp.fire()) {
        state := Mux(io.out.coh.resp.bits.isProbeHit(), s_probeForward, s_memReadReq)
      }
    }
    is (s_probeForward) {
      val thisResp = io.in.resp
      thisResp <> io.out.coh.resp
      when (thisResp.fire() && thisResp.bits.isReadLast()) { state := s_idle }
    }
    is (s_memReadReq) {
      io.out.mem.req.bits := reqLatch
      io.out.mem.req.valid := true.B
      when (io.out.mem.req.fire()) { state := s_memReadResp }
    }
    is (s_memReadResp) { when (io.out.mem.resp.fire() && io.out.mem.resp.bits.isReadLast()) { state := s_idle } }
    is (s_memWriteResp) { when (io.out.mem.resp.fire()) { state := s_idle } }
  }
}

class Prefetcher extends Module with HasPrefetcherParameter {
	val io = IO(new Bundle {
		val in = Flipped(Decoupled(new SimpleBusReqBundle))
		val out = Decoupled(new SimpleBusReqBundle)
	})
  val getNewReq = RegInit(false.B)
  val prefetchReq = RegNext(io.in.bits)
  prefetchReq.cmd := SimpleBusCmd.prefetch
  prefetchReq.addr := io.in.bits.addr + XLEN.U

  //lastReqAddr not be initted, in vivado simulation maybe fail
  //val lastReqAddr = (RegEnable(io.in.bits.addr, io.in.fire()))
  val lastReqAddr = RegInit(0.U(AddrBits.W))
  when (io.in.fire()) {
     lastReqAddr := io.in.bits.addr
  }
  val thisReqAddr = io.in.bits.addr
  val lineMask = Cat(Fill(AddrBits - 6, 1.U(1.W)), 0.U(6.W))
  val neqAddr = (thisReqAddr & lineMask) =/= (lastReqAddr & lineMask)

  when (!getNewReq) {
    io.out.bits <> io.in.bits
    io.out.valid := io.in.valid
    io.in.ready := !io.in.valid || io.out.fire()
    getNewReq := io.in.fire() && io.in.bits.isBurst() && neqAddr
  }.otherwise {
    io.out.bits <> prefetchReq
    io.out.valid := !AddressSpace.isMMIO(prefetchReq.addr)
    io.in.ready := false.B
    getNewReq := !(io.out.fire() || AddressSpace.isMMIO(prefetchReq.addr))
  }
}

case class CacheConfig (
  ro: Boolean = false,
  name: String = "cache",
  userBits: Int = 0,
  idBits: Int = 0,
  cacheLevel: Int = 1,

  totalSize: Int = 4, // Kbytes
  ways: Int = 4
)

class SingleLevelCacheSpecExternal extends Module {

  implicit val cacheConfig = new CacheConfig()
  val port = IO(new CacheIO())

  // Simplest case: data-cache, no coh manager, no dynamic exec 
  class CacheWrap(implicit val cacheConfig: CacheConfig) extends Module with HasNutCoreConst with HasCacheConst with HasCacheIO{        

    val s1 = Module(new CacheStage1)
    val s2 = Module(new CacheStage2)
    val s3 = Module(new CacheStage3)
    val metaArray = Module(new SRAMTemplateWithArbiter(nRead = 1, new MetaBundle, set = Sets, way = Ways))
    val dataArray = Module(new SRAMTemplateWithArbiter(nRead = 2, new DataBundle, set = Sets * LineBeats, way = Ways))

    val mem = Module(new AXI4RAM(memByte = 4096))

    // we are investigating the read request, thus forbid write request for this moment
    // since SRAMTemplateWithArbiter defines SRAMTemplate as single-ported
    // assume(!metaArray.io.w.req.valid && !dataArray.io.w.req.valid)
    // assume(!observe(metaArray.ram.resetState))

    val arb = Module(new Arbiter(new SimpleBusReqBundle(userBits = userBits, idBits = idBits), hasCohInt + 1))
    arb.io.in(hasCohInt + 0) <> io.in.req

    s1.io.in <> arb.io.out

    PipelineConnect(s1.io.out, s2.io.in, s2.io.out.fire, io.flush(0))
    PipelineConnect(s2.io.out, s3.io.in, s3.io.isFinish, io.flush(1))
    io.in.resp <> s3.io.out
    s3.io.flush := io.flush(1)

    // io.out.mem <> s3.io.mem
    mem.io.in <> s3.io.mem.toAXI4()
    io.out.mem <> DontCare

    io.mmio <> s3.io.mmio
    io.empty := !s2.io.in.valid && !s3.io.in.valid

    io.in.resp.valid := Mux(s3.io.out.valid && s3.io.out.bits.isPrefetch(), false.B, s3.io.out.valid || s3.io.dataReadRespToL1)

    // connecting to coh manager 
    val cohReq = io.out.coh.req.bits
    val coh = Wire(new SimpleBusReqBundle(userBits = userBits, idBits = idBits))
    coh.apply(addr = cohReq.addr, cmd = cohReq.cmd, size = cohReq.size, wdata = cohReq.wdata, wmask = cohReq.wmask)
    arb.io.in(0).bits := coh
    arb.io.in(0).valid := io.out.coh.req.valid
    io.out.coh.req.ready := arb.io.in(0).ready
    io.out.coh.resp <> s3.io.cohResp

    metaArray.io.r(0) <> s1.io.metaReadBus
    dataArray.io.r(0) <> s1.io.dataReadBus
    dataArray.io.r(1) <> s3.io.dataReadBus

    metaArray.io.w <> s3.io.metaWriteBus
    dataArray.io.w <> s3.io.dataWriteBus

    s2.io.metaReadResp := s1.io.metaReadBus.resp.data
    s2.io.dataReadResp := s1.io.dataReadBus.resp.data
    s2.io.dataWriteBus := s3.io.dataWriteBus
    s2.io.metaWriteBus := s3.io.metaWriteBus


    // ==================================================================================
    // FORMAL PROPERTIES START HERE

    // # 1. Memory Consistency (WEAK)
    //
    // We can describe the memory consistency in various ways. From behavioral perspective,
    // we can describe this property as:
    //
    // When a write request followed by a read request, if both requests are accessing same
    // address, then the read request always obtain the data in write request.
    //
    // Since we need to write formal property in assertion, we need to rephrase the 
    // property with the present tense.
    //
    // ==
    // If a write request WAS made, and NO MORE write request WAS made SINCE THEN, then
    // FOR ALL read request accesing same address SINCE THEN, the responses ARE the same
    // data as the data included in the write request.
    // ==
    // 
    // Note that "NO MORE write request WAS made since the last write" is a strong condition.
    // strong condition leads weak conclusion. And this condition doesn't cover all cache
    // behaviors since the consecutive save and load requests always cause read-hit, the most
    // simple case of cache behavior.
    //
    // # 2. Memory Consistency (STRONGER)
    //
    // A stronger version of this property is simply replace the "NO MORE ..." with "write
    // request with different address can be made". 
    //
    // And we wanted to examine the most complex situation. We are creating a SAVE-SAVE-LOAD
    // request triology, which stisfies;
    // 1. the 1st SAVE and final LOAD are accessing exact same address
    // 2. the 1st SAVE and 2nd SAVE are accessing different address with SAME DATA INDEX
    //
    // In this case, the first two saves will cause a write-miss, which further causes 
    // reallocation and write forward. the last save will cause another read-miss, which
    // causes two 

    // ::::XXX::::XXX::::XXX::::
    // Note about assumptions
    //
    // To be confirmed:
    // 1. is waymask in Stage 2 guaranteed to be One-Hot? or Mux1H will give undetermined
    //    result and cause wrong hit.
    //    Possible solution: use PriorityEncoder instead.
    //
    // 2. Is the register in PipelineConnect initialized? Unintialized register in the
    //    PipelineConnect will affect all the fields in a req.
    //
    // NOTE:
    // In the code below, the term "read" and "write" are replaced with "load" and "save"
    // for brevity and typographical beautiness.

    assume(io.in.req.bits.wdata === "hDECAF000DEADBEEF".U)
    assume(io.in.req.bits.wmask === "hFF".U)
    assume(io.in.req.bits.cmd === "b0001".U || io.in.req.bits.cmd === "b0000".U)
    assume(io.flush === 0.U)
    assume(!io.mmio.req.valid)
    assume(!io.out.coh.req.valid)

    // assert(PopCount(s2.io.out.bits.waymask) === 1.U)
    // assert(PopCount(s3.io.in.bits.waymask) === 1.U)

    val isSave       = io.in.req.valid && (
                       // io.in.req.bits.cmd === "b0001".U ||
                       // io.in.req.bits.cmd === "b0011".U ||
                       // io.in.req.bits.cmd === "b0111".U
                       io.in.req.bits.cmd === "b0001".U
                     )

    val isLoad       = io.in.req.valid && (
                       // io.in.req.bits.cmd === "b0000".U ||
                       // io.in.req.bits.cmd === "b0010".U ||
                       // io.in.req.bits.cmd === "b0100".U ||
                       // io.in.req.bits.cmd === "b1000".U
                       io.in.req.bits.cmd === "b0000".U  
                     )


    val isLoadOnce = RegInit(false.B)
    val isSaveOnce = RegInit(false.B)
    val isSubSaveReady = RegInit(false.B)
    val isSubSaveOnce = RegInit(false.B)

    val isSaveDone = io.in.resp.valid && io.in.resp.bits.isWriteResp()
    val isSaveDoneOnce = RegInit(false.B) 
    val isSubSaveDoneOnce = RegInit(false.B)

    val isLoadDone = io.in.resp.valid && (
      io.in.resp.bits.isReadLast()  || 
      io.in.resp.bits.isProbeHit()  ||
      io.in.resp.bits.isProbeMiss()
    )
    val isLoadDoneOnce = RegInit(false.B)

    // a counter for guarantee the sequence of save / load req.
    val (reqSeqCont, reqSeqWarp) = Counter(isLoad || isSave, 1048576)
    val saveSeq = RegEnable(reqSeqCont, 0.U, isSaveOnce && isSaveDone)
    val subSaveSeq = RegEnable(reqSeqCont, 0.U, isSubSaveOnce && isSaveDone)
    val loadSeq = RegEnable(reqSeqCont, 0.U, isLoadOnce && isLoadDone) 

    // During read or write request, the isLoadOnce / isSaveOnce is up.
    // when they are up, but isLoadDone or or isSaveDone, they are down.
    when (isSave) {
      //for the first save
      when (!isSubSaveReady) {
        isSaveOnce := true.B
        isSubSaveReady := true.B

      // when the first save is done
      }.elsewhen(!isSaveOnce) {
        isSubSaveOnce := true.B
      }
    }

    when (isSaveOnce && isSaveDone) {
      isSaveOnce := false.B
    }

    when (isSubSaveOnce && isSaveDone) {
      isSubSaveOnce := false.B
    }

    when (isLoad) {
      isLoadOnce := true.B
    }
    when (isLoadOnce && isLoadDone) {
      isLoadOnce := false.B
    }

    // new read / write request should not be made when the last one
    // is finishing, or the saved addr and data in register will be
    // overwritten
    assume(!(isLoad && isLoadDone))
    assume(!(isSave && isSaveDone))
    

    // when either kind of request is on-the-fly, no new requests
    // should be made.
    when (isSaveOnce || isLoadOnce || isSubSaveOnce) {
      assume (!io.in.req.valid)
    }

    // stored addr and data for comparison
    // these are interested signal for the formal property
    val loadAddrPast = RegEnable(io.in.req.bits.addr,  isLoad)
    val saveAddrPast = RegEnable(io.in.req.bits.addr,  isSave && !isSubSaveReady)
    val saveDataPast = RegEnable(io.in.req.bits.wdata, isSave && !isSubSaveReady)

    val subSaveAddrPast = RegEnable(io.in.req.bits.addr,  isSave && isSubSaveReady)
   
    // val s_idle :: s_memReadReq :: s_memReadResp :: s_memWriteReq :: s_memWriteResp :: s_mmioReq :: s_mmioResp :: s_wait_resp :: s_release :: Nil = Enum(9)

    // ========================================================================================
    // Excerpted from original code
    //
    // io.out.valid := io.in.valid && Mux(req.isBurst() && (cacheLevel == 2).B,
    //   Mux(req.isWriteSingle && (hit || !hit && state === s_wait_resp), true.B, (state === s_memReadResp && io.mem.resp.fire && req.cmd === SimpleBusCmd.readBurst)) || (respToL1Fire && respToL1Last && state === s_release),
    //   Mux(probe, false.B, Mux(hit, true.B, Mux(req.isWriteSingle || mmio, state === s_wait_resp, afterFirstRead && !alreadyOutFire)))
    // )
    //  io.isFinish := Mux(probe, io.cohResp.fire && Mux(miss, state === s_idle, (state === s_release) && releaseLast),
    //   Mux(hit || req.isWriteSingle, io.out.fire, (state === s_wait_resp) && (io.out.fire || alreadyOutFire))
    // )
    //
    //

    // assert(!(isSaveOnce && !isLoadOnce && isSaveDone))        // Expected to FAIL 
    
    // This won't work since isWrite doesn't detect other bits. We also leave
    // it to be proved that the MMU will ALWAYS send legal write command.
    // when (isSaveDone) { assert(isSaveOnce) }

    // This doesn't describe the order of arrival of two consecutive requests, since
    // isLoad can happen very early than isSave, yet isLoadDone still happen in the
    // current cycle, since the read request is long enough to include another write
    // request. Thus we have to use a counter to distinguish the order of two request
    //
    // when (isSaveOnce && isSaveDoneOnce && isLoadOnce && isLoadDone) {

    // assert(!(saveSeq > 0.U && loadSeq > 0.U && loadSeq > saveSeq))

    // ===================================================================================
    // 1. THE SAVE-LOAD CASE:
    //
    // Test if the simplest case is hold.  When 
    // 
    when (saveSeq > 0.U && loadSeq > 0.U) {

      when (loadSeq > saveSeq && isLoadOnce && isLoadDone) { 
        // assert(!(saveAddrPast === loadAddrPast))              // Expected to FAIL
        when (saveAddrPast === loadAddrPast){
          assert(saveDataPast === io.in.resp.bits.rdata)
        }
      }

    }

    when (saveSeq > 0.U && subSaveSeq > 0.U && loadSeq > 0.U) {

      when (loadSeq > subSaveSeq && subSaveSeq > saveSeq && isLoadOnce && isLoadDone) { 
        // assert(!(saveAddrPast === loadAddrPast))              // Expected to FAIL
        when (
          saveAddrPast === loadAddrPast && 
          saveAddrPast =/= subSaveAddrPast && 
          getDataIdx(saveAddrPast) === getDataIdx(subSaveAddrPast)) {
          assert(saveDataPast === io.in.resp.bits.rdata)
        }
      }

    }

    println("CACHE READ ONLY: " + ro)
    println("CACHE LEVEL: " + cacheLevel)
  }

  val mod = Module(new CacheWrap());
  mod.io <> port
}

class DoubleLevelCacheSpecExternal extends Module with HasNutCoreConst{

    val io = IO(new Bundle {
      val in = Flipped(new SimpleBusUC)
    })

    // preventing from boring with lsuMMIO
    val cacheDataOut = Cache(in = io.in, mmio = 0.U.asTypeOf(new SimpleBusUC) :: Nil, flush = "b00".U, empty = Wire(Bool()))(CacheConfig(totalSize = 1, name = "cacheData"))

    cacheDataOut.coh.resp.ready := true.B
    cacheDataOut.coh.req.valid  := false.B
    cacheDataOut.coh.req.bits := DontCare
    

    // val cohMan = Module(new CoherenceManager)
    val xbar = Module(new SimpleBusCrossbarNto1(2))
    // xbar.io.in(0) <> cohMan.io.out.mem
    xbar.io.in(0) <> DontCare 
    xbar.io.in(1) <> cacheDataOut.mem

    // preventing from using prefetch
    val cacheLevel2Out = Cache(in = xbar.io.out, mmio = 0.U.asTypeOf(new SimpleBusUC) :: Nil, flush = "b00".U, empty = Wire(Bool()))(CacheConfig(totalSize = 4, name = "cacheMem", cacheLevel = 2))

    cacheLevel2Out.coh.resp.ready := true.B
    cacheLevel2Out.coh.req.valid  := false.B
    cacheLevel2Out.coh.req.bits := DontCare
    
    val mem = Module(new AXI4RAM(memByte = 4096))

    mem.in <> cacheLevel2Out.mem.toAXI4()


    assume(io.in.req.bits.wdata === "hDECAF000DEADBEEF".U)
    assume(io.in.req.bits.wmask === "hFF".U)
    assume(io.in.req.bits.cmd === "b0001".U || io.in.req.bits.cmd === "b0000".U)

    val isSave       = io.in.req.valid && (
                       io.in.req.bits.cmd === "b0001".U
                     )

    val isLoad       = io.in.req.valid && (
                       io.in.req.bits.cmd === "b0000".U  
                     )


    val isLoadOnce = RegInit(false.B)
    val isSaveOnce = RegInit(false.B)

    val isSaveDone = io.in.resp.valid && io.in.resp.bits.isWriteResp()
    val isSaveDoneOnce = RegInit(false.B) 

    val isLoadDone = io.in.resp.valid && (
      io.in.resp.bits.isReadLast()  || 
      io.in.resp.bits.isProbeHit()  ||
      io.in.resp.bits.isProbeMiss()
    )
    val isLoadDoneOnce = RegInit(false.B)

    // a counter for guarantee the sequence of save / load req.
    val (reqSeqCont, reqSeqWarp) = Counter(isLoad || isSave, 1048576)
    val saveSeq = RegEnable(reqSeqCont, 0.U, isSaveOnce && isSaveDone)
    val loadSeq = RegEnable(reqSeqCont, 0.U, isLoadOnce && isLoadDone) 

    // During read or write request, the isLoadOnce / isSaveOnce is up.
    // when they are up, but isLoadDone or or isSaveDone, they are down.
    when (isSave) {
      isSaveOnce := true.B
    }
    when (isSaveOnce && isSaveDone) {
      isSaveOnce := false.B
    }

    when (isLoad) {
      isLoadOnce := true.B
    }
    when (isLoadOnce && isLoadDone) {
      isLoadOnce := false.B
    }

    // new read / write request should not be made when the last one
    // is finishing, or the saved addr and data in register will be
    // overwritten
    assume(!(isLoad && isLoadDone))
    assume(!(isSave && isSaveDone))

    // when either kind of request is on-the-fly, no new requests
    // should be made.
    when (isSaveOnce || isLoadOnce) {
      assume (!io.in.req.valid)
    }

    // stored addr and data for comparison
    // these are interested signal for the formal property
    val saveAddrPast = RegEnable(io.in.req.bits.addr,  isSave)
    val saveDataPast = RegEnable(io.in.req.bits.wdata, isSave)
    val loadAddrPast = RegEnable(io.in.req.bits.addr,  isLoad)
      // assert(!(saveSeq > 0.U && loadSeq > 0.U && loadSeq > saveSeq))
    when (saveSeq > 0.U && loadSeq > 0.U && loadSeq > saveSeq && isLoadOnce && isLoadDone) {
      // assert(!(saveAddrPast === loadAddrPast))              // Expected to FAIL
      when (saveAddrPast === loadAddrPast) {
        assert(saveDataPast === io.in.resp.bits.rdata)
      }
    }


}
